const std = @import("std");
const Segment = @import("VMWriter.zig").Segment;
const SymbolTable = @This();

pub const Kind = enum {
    field,
    static,
    argument,
    local,

    pub fn toSegment(kind: Kind) Segment {
        return switch (kind) {
            .field => .this,
            inline else => |k| @field(Segment, @tagName(k)),
        };
    }
};

const VarAttr = struct {
    type: []const u8,
    kind: Kind,
    index: usize,
};

indexes: [4]usize = @splat(0),
table: struct {
    class: std.StringHashMap(VarAttr),
    subroutine: std.StringHashMap(VarAttr),
},

pub fn init(allocator: std.mem.Allocator) SymbolTable {
    return .{
        .table = .{
            .class = .init(allocator),
            .subroutine = .init(allocator),
        },
    };
}

pub fn deinit(self: *SymbolTable) void {
    self.table.class.deinit();
    self.table.subroutine.deinit();
}

pub fn reset(self: *SymbolTable) void {
    self.table.subroutine.clearRetainingCapacity();
    self.indexes[@intFromEnum(Kind.argument)] = 0;
    self.indexes[@intFromEnum(Kind.local)] = 0;
}

pub fn define(self: *SymbolTable, name: []const u8, @"type": []const u8, kind: Kind) !void {
    const kind_index = @intFromEnum(kind);
    var table = if (kind_index <= @intFromEnum(Kind.static)) &self.table.class else &self.table.subroutine;

    try table.put(name, .{ .type = @"type", .kind = kind, .index = self.indexes[kind_index] });
    self.indexes[kind_index] += 1;
}

pub fn varCount(self: SymbolTable, kind: Kind) usize {
    return self.indexes[@intFromEnum(kind)];
}

pub fn kindOf(self: SymbolTable, name: []const u8) ?Segment {
    const kind: Kind = if (self.table.subroutine.get(name)) |variable| variable.kind else blk: {
        break :blk if (self.table.class.get(name)) |v| v.kind else return null;
    };
    return kind.toSegment();
}

pub fn typeOf(self: SymbolTable, name: []const u8) []const u8 {
    return if (self.table.subroutine.get(name)) |variable| variable.type else self.table.class.get(name).?.type;
}

pub fn indexOf(self: SymbolTable, name: []const u8) usize {
    return if (self.table.subroutine.get(name)) |variable| variable.index else self.table.class.get(name).?.index;
}
